import pandas as pd

from ruleng.match_rule import Analytics

"""
class Analytics:
    def __init__(self, rule=None, rule_id=-1, host='192.168.1.143',
                 redis_host="localhost", redis_port=6379,
                 search_index='btgd_xml*', size=100000, time_out=100,
                 precision=4000, verbose=0, num_worker=10, record_size=1,
                 default_time_interval='3M', steps=1000, num_records=100000,
                 violation_limit=10000, mode=0, thread_mode=0,
                 log_file_name='', speedup_mode=2,
                 max_container_print_size=20,
                 get_timed_assoc_tab_through_aggs=True,
                 exception_retries=3, bailout_2nd_pass_threshold=0.5,
                 finished_time_spec=[],
                 scroll_keep_alive='1m',
                 multi_level_aggs=True,
                 doc_time_field='origin_time',
                 doc_time_hms_field='',
                 doc_timestamp_field='@timestamp',
                 doc_type='',
                 redis_passwd='admin'
                 ):
"""


class AnalyticsPandas(Analytics):
    def get_recent_scroll_records_sources(self, one_match_res, **kwargs):
        for d in self.get_recent_scroll_records(one_match_res, **kwargs):
            yield d["_source"]

    def get_recent_scroll_records_df(self, one_match_res, **kwargs):
        return pd.DataFrame(
            self.get_recent_scroll_records_sources(one_match_res, **kwargs)
        )

    def get_df_by_field(
        self,
        field_name,
        field_vals,
        redis_off=True,
        doc_time_field=None,
        search_index=None,
        start_date="1990-01-01 00:00:00",
        end_date="2030-12-31 23:00:00",
        **kwargs
    ):
        assert isinstance(field_name, str) and isinstance(field_vals, list)
        self.redis_off = redis_off
        if doc_time_field:
            self.doc_time_field = doc_time_field
        if search_index:
            self.search_index = search_index
        self.rule = [[field_name, "query", "in", field_vals]]
        res = self.vector_rule([start_date, end_date])
        if len(res) < 1:
            return pd.DataFrame()
        return self.get_recent_scroll_records_df(
            one_match_res=res[0], **kwargs
        ).astype({self.doc_time_field: "datetime64[ns]"})

    def get_df_by_rule(
        self,
        rule,
        search_index=None,
        doc_time_field=None,
        redis_off=True,
        start_date="1990-01-01 00:00:00",
        end_date="2030-12-31 23:00:00",
        **kwargs
    ):
        self.redis_off = redis_off
        if doc_time_field:
            self.doc_time_field = doc_time_field
        if search_index:
            self.search_index = search_index
        self.rule = rule
        res = self.vector_rule([start_date, end_date])
        if len(res) < 1:
            return pd.DataFrame()
        return self.get_recent_scroll_records_df(
            one_match_res=res[0], **kwargs
        ).astype({self.doc_time_field: "datetime64[ns]"})


if __name__ == "__main__":
    es_host = "192.168.83.74:9200"
    search_index = "split_res-31403"
    field_name = "ORG_NO"
    field_vals = ["3140303"]
    doc_time_field = "CREATE_DATE"
    ap = AnalyticsPandas(host=es_host)
    df = ap.get_df_by_field(
        field_name,
        field_vals,
        doc_time_field=doc_time_field,
        search_index=search_index,
    )
    print(df)
    print(df.columns)
    print(df.dtypes)
